#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Net4Me - 异地组网解决方案主程序
基于SSH反向隧道技术实现内网穿透
"""

import json
import logging
import argparse
import time
import threading
import paramiko
from pathlib import Path


class SSHTunnelManager:
    """SSH隧道管理器"""
    
    def __init__(self, config_path: str = "config.json"):
        """
        初始化SSH隧道管理器
        
        Args:
            config_path (str): 配置文件路径
        """
        self.config_path = config_path
        self.config = self._load_config()
        self.ssh_client = None
        self.transport = None
        self.is_connected = False
        self._setup_logging()
        
    def _load_config(self) -> dict:
        """
        加载配置文件
        
        Returns:
            dict: 配置信息字典
        """
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except FileNotFoundError:
            logging.error(f"配置文件 {self.config_path} 未找到")
            raise
        except json.JSONDecodeError:
            logging.error(f"配置文件 {self.config_path} 格式错误")
            raise
            
    def _setup_logging(self):
        """设置日志记录"""
        log_level = getattr(logging, self.config['tunnel']['log_level'], logging.INFO)
        logging.basicConfig(
            level=log_level,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler('net4me.log'),
                logging.StreamHandler()
            ]
        )
        
    def _create_ssh_client(self) -> paramiko.SSHClient:
        """
        创建SSH客户端
        
        Returns:
            paramiko.SSHClient: SSH客户端实例
        """
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        return ssh
        
    def connect(self) -> bool:
        """
        建立SSH连接并创建反向隧道
        
        Returns:
            bool: 连接是否成功
        """
        try:
            # 获取远程服务器配置
            remote_config = self.config['remote_server']
            
            # 创建SSH客户端
            self.ssh_client = self._create_ssh_client()
            
            # 连接参数
            connect_kwargs = {
                'hostname': remote_config['host'],
                'port': remote_config['port'],
                'username': remote_config['username']
            }
            
            # 根据配置选择认证方式
            if remote_config['private_key_path']:
                # 使用私钥认证
                private_key_path = Path(remote_config['private_key_path'])
                if private_key_path.exists():
                    private_key = paramiko.RSAKey.from_private_key_file(str(private_key_path))
                    connect_kwargs['pkey'] = private_key
                else:
                    logging.warning(f"私钥文件 {private_key_path} 不存在，尝试使用密码认证")
                    if remote_config['password']:
                        connect_kwargs['password'] = remote_config['password']
                    else:
                        raise ValueError("未提供有效的认证信息")
            else:
                # 使用密码认证
                if remote_config['password']:
                    connect_kwargs['password'] = remote_config['password']
                else:
                    raise ValueError("未提供有效的认证信息")
            
            # 建立SSH连接
            logging.info(f"正在连接到远程服务器 {remote_config['host']}:{remote_config['port']}")
            self.ssh_client.connect(**connect_kwargs)
            self.transport = self.ssh_client.get_transport()
            
            # 设置keepalive
            keepalive_interval = self.config['tunnel']['keepalive_interval']
            self.transport.set_keepalive(keepalive_interval)
            logging.info(f"已设置keepalive，间隔 {keepalive_interval} 秒")
            
            # 请求反向隧道
            remote_port = remote_config['remote_port']
            local_host = self.config['local_machine']['host']
            local_port = self.config['local_machine']['port']
            
            logging.info(f"正在创建反向隧道: 远程端口 {remote_port} -> 本地 {local_host}:{local_port}")
            self.transport.request_port_forward('', remote_port, handler=self._forward_handler)
            
            self.is_connected = True
            logging.info("SSH连接和反向隧道已建立")
            return True
            
        except Exception as e:
            logging.error(f"连接失败: {str(e)}")
            self.is_connected = False
            return False
            
    def _forward_handler(self, chan, origin, server):
        """
        反向隧道处理器
        
        Args:
            chan: SSH通道
            origin: 连接来源
            server: 服务器信息
        """
        try:
            logging.info(f"收到连接请求来自 {origin[0]}:{origin[1]}")
            
            # 连接到本地SSH服务
            local_config = self.config['local_machine']
            local_sock = paramiko.Transport((local_config['host'], local_config['port']))
            
            # 本地认证
            if local_config['private_key_path']:
                private_key = paramiko.RSAKey.from_private_key_file(local_config['private_key_path'])
                local_sock.connect(username=local_config['username'], pkey=private_key)
            else:
                local_sock.connect(username=local_config['username'], password=local_config['password'])
                
            # 创建转发通道
            local_chan = local_sock.open_session()
            local_chan.settimeout(60)
            local_chan.get_pty()
            local_chan.invoke_shell()
            
            # 建立双向转发
            def forward_out():
                try:
                    while True:
                        data = chan.recv(1024)
                        if not data:
                            break
                        local_chan.send(data)
                except:
                    pass
                
            def forward_in():
                try:
                    while True:
                        data = local_chan.recv(1024)
                        if not data:
                            break
                        chan.send(data)
                except:
                    pass
                    
            threading.Thread(target=forward_out, daemon=True).start()
            threading.Thread(target=forward_in, daemon=True).start()
            
        except Exception as e:
            logging.error(f"处理连接请求时出错: {str(e)}")
            try:
                chan.close()
            except:
                pass
                
    def disconnect(self):
        """断开SSH连接"""
        try:
            if self.transport:
                self.transport.close()
            if self.ssh_client:
                self.ssh_client.close()
            self.is_connected = False
            logging.info("SSH连接已断开")
        except Exception as e:
            logging.error(f"断开连接时出错: {str(e)}")
            
    def reconnect(self):
        """
        重新连接
        
        Returns:
            bool: 重连是否成功
        """
        logging.info("尝试重新连接...")
        self.disconnect()
        return self.connect()
        
    def run(self):
        """运行主循环"""
        logging.info("Net4Me客户端启动")
        
        # 首次连接
        if not self.connect():
            logging.error("首次连接失败")
            return
            
        # 保持运行并监控连接状态
        reconnect_interval = self.config['tunnel']['reconnect_interval']
        try:
            while True:
                if not self.is_connected:
                    logging.warning("检测到连接断开，尝试重新连接...")
                    if not self.reconnect():
                        logging.error(f"重连失败，{reconnect_interval}秒后重试")
                        time.sleep(reconnect_interval)
                        continue
                        
                # 检查连接状态
                if self.transport and not self.transport.is_active():
                    logging.warning("检测到传输通道不活跃，重新连接...")
                    if not self.reconnect():
                        logging.error(f"重连失败，{reconnect_interval}秒后重试")
                        time.sleep(reconnect_interval)
                        continue
                        
                time.sleep(10)  # 每10秒检查一次
                
        except KeyboardInterrupt:
            logging.info("收到中断信号，正在关闭连接...")
        finally:
            self.disconnect()


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='Net4Me - 异地组网解决方案')
    parser.add_argument('--mode', choices=['client'], default='client', 
                       help='运行模式 (目前仅支持client模式)')
    parser.add_argument('--config', default='config.json', 
                       help='配置文件路径')
    
    args = parser.parse_args()
    
    try:
        tunnel_manager = SSHTunnelManager(args.config)
        tunnel_manager.run()
    except Exception as e:
        logging.error(f"程序运行出错: {str(e)}")
        return 1
        
    return 0


if __name__ == "__main__":
    exit(main())